import os
import pprint
from glob import glob

import nibabel
import numpy as np
import torch
import torchio
from tqdm import tqdm

from config.KiPA2022_config import ModifiedUNet3DConfigure
from models.EvaluationMetrics import Dice
from utils.DataUtils import get_patch_without_mask
from utils.Utils import get_size_from_indexlist


def do_evaluation(eval_configure, model_configure_list):
    """
    This function is used to predict the evaluation data which don't have the label
    :param eval_configure:
    :param model_configure_list:
    :return:
    """

    os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
    device = torch.device('cuda')
    compute_dice = Dice(model_configure_list[0].class_num)
    sample_list = glob(os.path.join(eval_configure.valid_data_root_path, "*"))

    normal_transform = torchio.ZNormalization()

    # Box_boundary 中坐标前开后闭
    for sample_path in sample_list:
        with torch.no_grad():
            if predict_configure.run_type == "train":
                sample_name = os.path.basename(sample_path)
                label_data_path = os.path.join(sample_path, "label", sample_name + ".nii.gz")
                data_affine = nibabel.load(label_data_path).affine
                valid_label = torch.from_numpy(nibabel.load(label_data_path).get_data()).unsqueeze(0).unsqueeze(0)
                predict_dir_path = os.path.join(eval_configure.valid_data_root_path,
                                                sample_name, "label")
            else:
                data_affine = nibabel.load(sample_path).affine
                sample_name = os.path.basename(sample_path).split(".")[0]
                predict_dir_path = os.path.join(eval_configure.valid_data_root_path, "predict")
            if not os.path.exists(predict_dir_path):
                os.mkdir(predict_dir_path)
            model_predict_result_list = []
            for configure in model_configure_list:
                model = configure.model
                model_name = configure.model.__str__().split("\n")[0].replace("(", "")
                if predict_configure.run_type == "train":
                    predict_file_path = os.path.join(eval_configure.predict_file_path, sample_name, "label")
                    predict_file_path = os.path.join(predict_file_path,
                                                     sample_name + "_valid_" + model_name + ".nii.gz")
                else:
                    predict_file_path = os.path.join(predict_dir_path, sample_name + ".nii.gz")
                weights_path = configure.weights_path
                check_point = torch.load(weights_path)
                model.load_state_dict(check_point)
                model = model.to(device)
                model.eval()
                valid_data_list, valid_label_list, index_list = get_patch_without_mask(sample_path,
                                                                                       configure.patch_size, 2,
                                                                                       eval_configure.dataset_name,
                                                                                       eval_configure.run_type)
                model_predict_result = torch.zeros(
                    [1, predict_configure.class_num] + get_size_from_indexlist(index_list))
                for valid_data_patch, valid_label_patch, index in tqdm(
                        zip(valid_data_list, valid_label_list, index_list)):
                    valid_data_patch = torch.from_numpy(valid_data_patch).to(device, dtype=torch.float32)
                    valid_data_patch = normal_transform(valid_data_patch.unsqueeze(0))
                    valid_data_patch = valid_data_patch.unsqueeze(0)

                    # active_predict, predict = model(test_patch)
                    patch_predict_result = model(valid_data_patch)

                    model_predict_result[0, :, index[0]:index[1], index[2]:index[3],
                    index[4]:index[5]] = patch_predict_result.cpu()
                model_predict_result = torch.where(model_predict_result > configure.threshold, torch.tensor(1),
                                                   torch.tensor(0))
                if eval_configure.run_type == "train":
                    total_dice, dice_list = compute_dice(model_predict_result, valid_label)
                    pprint.pprint(dice_list)
                model_predict_result = torch.argmax(model_predict_result, dim=1, keepdim=True)

                model_predict_result_list.append(model_predict_result.numpy().astype(np.int8)[0][0])
            for predict_result in model_predict_result_list:
                nibabel.Nifti1Image(predict_result, data_affine).to_filename(predict_file_path)


if __name__ == '__main__':
    # class Modified3DUNetPredictConf:
    #     class_num = 1
    #     if os.environ.get("COMPUTERNAME") == "DESKTOP-16QO4BT":
    #         weights_path = r"H:\laboratory\project\blood_vessel\MICCAI2022-Draft\weights\MUNet_epoch_39"
    #     else:
    #         weights_path = r"C:\gs\code\parse2022-draft\weights\Modified3DUNet(1653108670972311\epoch_1"
    #     model = Modified3DUNet(1, class_num)
    #
    #     if os.environ.get("COMPUTERNAME") == "DESKTOP-16QO4BT":
    #         test_img_path = r"D:\dataset\vessel\Parse_2022_train_data\PA000144"
    #     else:
    #         test_img_path = r"C:\dataset\vessel\Parse_2022_train_data\PA000194"
    #     threshold = 0.95
    class predict_configure:
        class_num = 5
        run_type = "valid"
        dataset_name = "KiPA2022"
        valid_data_root_path = r"C:\dataset\other\open_image\image"
        predict_file_path = r"C:\dataset\other\open_image\image"


    class ModifiedUNet3DPredictConf(ModifiedUNet3DConfigure):
        class_num = 5
        weights_path = r"..\weights\KiPA2022_Modified3DUNet_Aug04-00-09-51\epoch_240global_step_26881"

        threshold = 0.5


    do_evaluation(predict_configure, [ModifiedUNet3DPredictConf])
